package raft

//
// this is an outline of the API that raft must expose to
// the service (or tester). see comments below for
// each of these functions for more details.
//
// rf = Make(...)
//   create a new Raft server.
// rf.Start(command interface{}) (index, term, isleader)
//   start agreement on a new log entry
// rf.GetState() (term, isLeader)
//   ask a Raft for its current term, and whether it thinks it is leader
// ApplyMsg
//   each time a new entry is committed to the log, each Raft peer
//   should send an ApplyMsg to the service (or tester)
//   in the same server.
//

import (
	"bytes"
	"fmt"
	"sync"
	"time"
)
import "sync/atomic"
import "../labrpc"
import "../labgob"
import "math/rand"

// ApplyMsg
// as each Raft peer becomes aware that successive log entries are
// committed, the peer should send an ApplyMsg to the service (or
// tester) on the same server, via the applyCh passed to Make(). set
// CommandValid to true to indicate that the ApplyMsg contains a newly
// committed log entry.
//
// in Lab 3 you'll want to send other kinds of messages (e.g.,
// snapshots) on the applyCh; at that point you can add fields to
// ApplyMsg, but set CommandValid to false for these other uses.
//
type ApplyMsg struct {
	CommandValid bool
	Command      interface{}
	CommandIndex int
}

type LogEntry struct {
	Command interface{}
	Term    int
	Index   int
}

type ServerRole int

const (
	Follower  ServerRole = 1
	Leader    ServerRole = 2
	Candidate ServerRole = 3
)

// Raft
// A Go object implementing a single Raft peer.
//
type Raft struct {
	mu        sync.Mutex          // Lock to protect shared access to this peer's state
	peers     []*labrpc.ClientEnd // RPC end points of all peers
	persister *Persister          // Object to hold this peer's persisted state
	me        int                 // this peer's index into peers[]
	dead      int32               // set by Kill()

	// Your data here (2A, 2B, 2C).
	// Look at the paper's Figure 2 for a description of what
	// state a Raft server must maintain.

	// persistent state
	currentTerm int
	votedFor    int
	log         []*LogEntry

	// volatile state
	commitIndex int
	lastApplied int

	// volatile state on leaders
	nextIndex  []int
	matchIndex []int

	// others
	role  ServerRole // follower -- 1, leader -- 2, candidate -- 3
	votes int        // number of votes of Candidate

	// channels between servers
	ApplyCh       chan ApplyMsg
	GrantVoteCh   chan bool
	WinElectionCh chan bool
	HeartBeatCh   chan bool
}

// GetState return currentTerm and whether this server
// believes it is the leader.
func (rf *Raft) GetState() (int, bool) {

	var term int
	var isLeader bool
	// Your code here (2A).
	term = rf.currentTerm
	isLeader = rf.role == Leader
	return term, isLeader
}

//
// save Raft's persistent state to stable storage,
// where it can later be retrieved after a crash and restart.
// see paper's Figure 2 for a description of what should be persistent.
//
func (rf *Raft) persist() {
	// Your code here (2C).
	// Example:
	w := new(bytes.Buffer)
	e := labgob.NewEncoder(w)
	e.Encode(rf.currentTerm)
	e.Encode(rf.votedFor)
	e.Encode(rf.log)
	data := w.Bytes()
	rf.persister.SaveRaftState(data)
}

//
// restore previously persisted state.
//
func (rf *Raft) readPersist(data []byte) {
	if data == nil || len(data) < 1 { // bootstrap without any state?
		return
	}
	// Your code here (2C).
	// Example:
	r := bytes.NewBuffer(data)
	d := labgob.NewDecoder(r)
	d.Decode(&rf.currentTerm)
	d.Decode(&rf.votedFor)
	d.Decode(&rf.log)
}

// RequestVoteArgs
// example RequestVote RPC arguments structure.
// field names must start with capital letters!
//
type RequestVoteArgs struct {
	// Your data here (2A, 2B).
	Term         int
	CandidateId  int
	LastLogIndex int
	LastLogTerm  int
}

// RequestVoteReply
// example RequestVote RPC reply structure.
// field names must start with capital letters!
//
type RequestVoteReply struct {
	// Your data here (2A).
	Term        int
	VoteGranted bool
}

// RequestVote
// example RequestVote RPC handler.
//
func (rf *Raft) RequestVote(args *RequestVoteArgs, reply *RequestVoteReply) {
	// Your code here (2A, 2B).
	rf.mu.Lock()
	defer rf.mu.Unlock()
	defer rf.persist()

	if rf.currentTerm > args.Term {
		// reject request with a stale term number
		reply.VoteGranted = false
		reply.Term = rf.currentTerm
		return
	}
	if rf.currentTerm < args.Term {
		// become follower and update current term
		rf.role = Follower
		rf.votedFor = -1
		rf.currentTerm = args.Term
	}
	reply.Term = rf.currentTerm
	reply.VoteGranted = false

	if (rf.votedFor == -1 || rf.votedFor == args.CandidateId) && rf.isUpToDate(args.LastLogTerm, args.LastLogIndex) {
		// vote for the candidate
		reply.VoteGranted = true
		rf.votedFor = args.CandidateId
		rf.GrantVoteCh <- true
	}
}

func (rf *Raft) isUpToDate(candidateTerm, candidateIndex int) bool {
	if candidateTerm == rf.lastLogTerm() {
		// here must be >=
		return candidateIndex >= rf.lastLogIndex()
	}
	return candidateTerm > rf.lastLogTerm()
}

//
// example code to send a RequestVote RPC to a server.
// server is the index of the target server in rf.peers[].
// expects RPC arguments in args.
// fills in *reply with RPC reply, so caller should
// pass &reply.
// the types of the args and reply passed to Call() must be
// the same as the types of the arguments declared in the
// handler function (including whether they are pointers).
//
// The labrpc package simulates a lossy network, in which servers
// may be unreachable, and in which requests and replies may be lost.
// Call() sends a request and waits for a reply. If a reply arrives
// within a timeout interval, Call() returns true; otherwise
// Call() returns false. Thus Call() may not return for a while.
// A false return can be caused by a dead server, a live server that
// can't be reached, a lost request, or a lost reply.
//
// Call() is guaranteed to return (perhaps after a delay) *except* if the
// handler function on the server side does not return.  Thus there
// is no need to implement your own timeouts around Call().
//
// look at the comments in ../labrpc/labrpc.go for more details.
//
// if you're having trouble getting RPC to work, check that you've
// capitalized all field names in structs passed over RPC, and
// that the caller passes the address of the reply struct with &, not
// the struct itself.
//
func (rf *Raft) sendRequestVote(server int, args *RequestVoteArgs, reqReply *RequestVoteReply) bool {
	ok := rf.peers[server].Call("Raft.RequestVote", args, reqReply)
	rf.mu.Lock()
	defer rf.mu.Unlock()
	defer rf.persist()

	if ok {
		if rf.role != Candidate || rf.currentTerm != args.Term {
			// invalid request
			return ok
		}
		// bigger term, turn into follower
		if reqReply.Term > rf.currentTerm {
			rf.debugLogger("candidate meet larger term, turn into follower")
			rf.role = Follower
			rf.currentTerm = reqReply.Term
			rf.votedFor = -1
			return ok
		}
		if reqReply.VoteGranted {
			rf.votes++
			if rf.votes > len(rf.peers)/2 {
				rf.debugLogger("candidate get enough votes, turn into leader")
				rf.role = Leader
				rf.persist()
				rf.nextIndex = make([]int, len(rf.peers))
				rf.matchIndex = make([]int, len(rf.peers))
				var nextIndexDefault = rf.lastLogIndex() + 1
				for server := range rf.nextIndex {
					rf.nextIndex[server] = nextIndexDefault
				}
				rf.WinElectionCh <- true
				return ok
			}
		}
	}

	return ok
}

type AppendEntriesArgs struct {
	Term         int
	LeaderId     int
	PrevLogIndex int
	PrevLogTerm  int
	Entries      []*LogEntry
	LeaderCommit int
}

type AppendEntriesReply struct {
	Term         int
	Success      bool
	NextTryIndex int
}

func (rf *Raft) AppendEntries(args *AppendEntriesArgs, reply *AppendEntriesReply) {
	rf.mu.Lock()
	defer rf.mu.Unlock()
	defer rf.persist()

	rf.debugLogger("Receiving AppendEntries from leader")
	reply.Success = false

	if args.Term < rf.currentTerm {
		// leader with lower term, reject
		reply.Term = rf.currentTerm
		reply.NextTryIndex = rf.lastLogIndex() + 1
		return
	}

	if args.Term > rf.currentTerm {
		// meet larger term, turn into follower
		rf.role = Follower
		rf.currentTerm = args.Term
		rf.votedFor = -1
	}

	// refresh timeout
	rf.HeartBeatCh <- true

	reply.Term = rf.currentTerm

	// if no entry at prevLogIndex, reply false (missing logs)
	if args.PrevLogIndex > rf.lastLogIndex() {
		reply.NextTryIndex = rf.lastLogIndex() + 1
		return
	}

	// Reply false if log doesn't contain an entry at prevLogIndex whose term matches prevLogTerm (§5.3)
	if args.PrevLogIndex >= 0 && rf.log[args.PrevLogIndex].Term != args.PrevLogTerm {
		// get NextTryIndex
		term := rf.log[args.PrevLogIndex].Term
		for i := args.PrevLogIndex - 1; i >= 0; i-- {
			if rf.log[i].Term != term {
				reply.NextTryIndex = i + 1
				break
			}
		}
	} else if args.PrevLogIndex >= -1 {
		// If an existing entry conflicts with a new one (same index but different terms), delete the existing entry and all that follow it (§5.3)
		rf.log = rf.log[:args.PrevLogIndex+1] // if prevLogIndex is -1, this deletion is also ok
		// Append any new entries not already in the log
		rf.debugLogger2B(fmt.Sprintf("appending logs, log: %v, leaderlog: %v", rf.log, args.Entries))
		rf.log = append(rf.log, args.Entries...)

		reply.Success = true
		reply.NextTryIndex = args.PrevLogIndex + len(args.Entries)

		// If leaderCommit > commitIndex, set commitIndex = min(leaderCommit, index of last new entry)
		if args.LeaderCommit > rf.commitIndex {
			rf.commitIndex = minInt(args.LeaderCommit, rf.lastLogIndex())
			go rf.applyLogs()
		}
	}
}

func minInt(x, y int) int {
	if x < y {
		return x
	}
	return y
}

func (rf *Raft) sendAppendEntries(server int, args *AppendEntriesArgs, reply *AppendEntriesReply) bool {
	ok := rf.peers[server].Call("Raft.AppendEntries", args, reply)
	rf.mu.Lock()
	defer rf.mu.Unlock()

	if !ok || rf.role != Leader || args.Term != rf.currentTerm {
		// invalid request
		return ok
	}
	if reply.Term > rf.currentTerm {
		// the leader faces a bigger term
		rf.debugLogger("Leader meet larger term, turn into follower")
		rf.debugLogger2B(fmt.Sprintf("leaderHandleAE: term changing from %v to %v", rf.currentTerm, reply.Term))
		rf.currentTerm = reply.Term
		rf.role = Follower
		rf.votedFor = -1
		rf.persist()
		return ok
	}
	if reply.Success {
		// AppendEntries RPC is successful
		// update nextIndex and matchIndex
		if len(args.Entries) > 0 {
			// not empty heartbeats, it contains log entries
			rf.nextIndex[server] = args.Entries[len(args.Entries)-1].Index + 1
			rf.matchIndex[server] = rf.nextIndex[server] - 1
		}
	} else {
		// AppendEntries RPC is unsuccessful because of log inconsistency
		// decrement nextIndex
		// using trick
		rf.nextIndex[server] = minInt(reply.NextTryIndex, rf.lastLogIndex())
	}

	// now handle commit index
	for N := rf.lastLogIndex(); N > rf.commitIndex && rf.log[N].Term == rf.currentTerm; N-- {
		count := 1
		for i := range rf.peers {
			if i != rf.me && rf.matchIndex[i] >= N {
				count++
			}
		}
		if count > len(rf.peers)/2 {
			rf.commitIndex = N
			go rf.applyLogs()
			break
		}
	}

	return ok
}

// Start
// the service using Raft (e.g. a k/v server) wants to start
// agreement on the next command to be appended to Raft's log. if this
// server isn't the leader, returns false. otherwise start the
// agreement and return immediately. there is no guarantee that this
// command will ever be committed to the Raft log, since the leader
// may fail or lose an election. even if the Raft instance has been killed,
// this function should return gracefully.
//
// the first return value is the index that the command will appear at
// if it's ever committed. the second return value is the current
// term. the third return value is true if this server believes it is
// the leader.
//
func (rf *Raft) Start(command interface{}) (int, int, bool) {
	index := -1
	term := -1
	isLeader := false

	// Your code here (2B).
	rf.mu.Lock()
	defer rf.mu.Unlock()

	rf.debugLogger2B(fmt.Sprintf("Start: command: %v, server: %v, isLeader: %v", command, rf.me, rf.role == Leader))
	isLeader = rf.role == Leader
	if !isLeader {
		// if this server isn't the leader, returns false
		return index, term, isLeader
	} else {
		// here this server is leader
		term = rf.currentTerm
		index = rf.lastLogIndex() + 1
		newLogEntry := LogEntry{Term: rf.currentTerm, Command: command, Index: index}
		rf.log = append(rf.log, &newLogEntry)
		rf.persist()
		return index, term, isLeader
	}
}

// Kill
// the tester doesn't halt goroutines created by Raft after each test,
// but it does call the Kill() method. your code can use killed() to
// check whether Kill() has been called. the use of atomic avoids the
// need for a lock.
//
// the issue is that long-running goroutines use memory and may chew
// up CPU time, perhaps causing later tests to fail and generating
// confusing debug output. any goroutine with a long-running loop
// should call killed() to check whether it should stop.
//
func (rf *Raft) Kill() {
	atomic.StoreInt32(&rf.dead, 1)
	// Your code here, if desired.
}

func (rf *Raft) killed() bool {
	z := atomic.LoadInt32(&rf.dead)
	return z == 1
}

// Make
// the service or tester wants to create a Raft server. the ports
// of all the Raft servers (including this one) are in peers[]. this
// server's port is peers[me]. all the servers' peers[] arrays
// have the same order. persister is a place for this server to
// save its persistent state, and also initially holds the most
// recent saved state, if any. applyCh is a channel on which the
// tester or service expects Raft to send ApplyMsg messages.
// Make() must return quickly, so it should start goroutines
// for any long-running work.
//
func Make(peers []*labrpc.ClientEnd, me int,
	persister *Persister, applyCh chan ApplyMsg) *Raft {
	rf := &Raft{}
	rf.peers = peers
	rf.persister = persister
	rf.me = me

	// Your initialization code here (2A, 2B, 2C).
	rf.role = Follower
	rf.votedFor = -1

	rf.log = append(rf.log, &LogEntry{
		Term: 0,
	})

	rand.Seed(time.Now().Unix())

	rf.ApplyCh = applyCh
	rf.HeartBeatCh = make(chan bool)
	rf.GrantVoteCh = make(chan bool)
	rf.WinElectionCh = make(chan bool)

	// initialize from state persisted before a crash
	rf.readPersist(persister.ReadRaftState())

	rf.persist()

	go rf.watchAction()

	return rf
}

func (rf *Raft) watchAction() {
	for !rf.killed() {
		switch rf.role {
		case Follower:
			select {
			case <-rf.GrantVoteCh:
				// grant vote for a certain candidate, begin another loop
			case <-rf.HeartBeatCh:
				// receiving heartbeat, begin another loop
			case <-time.After(time.Millisecond * time.Duration(rand.Intn(300)+200)):
				// election timeout
				rf.role = Candidate
				rf.persist()
			}
		case Leader:
			go rf.leaderAppendEntries()
			time.Sleep(time.Millisecond * 60)
		case Candidate:
			rf.mu.Lock()
			rf.currentTerm++
			rf.votes = 1
			rf.votedFor = rf.me
			rf.persist()
			rf.mu.Unlock()
			go rf.candidateRequestVotes()

			select {
			case <-rf.HeartBeatCh:
				// receiving heartbeat, turn into Follower
				rf.role = Follower
				rf.persist()
			case <-rf.WinElectionCh:
				// here the Candidate has already won the election, jump to the next loop, where the server is already the Leader
			case <-time.After(time.Millisecond * time.Duration(rand.Intn(300)+200)):
				// wait for election timeout
			}
		}
	}
}

func (rf *Raft) candidateRequestVotes() {
	rf.mu.Lock()
	reqArg := RequestVoteArgs{
		Term:         rf.currentTerm,
		CandidateId:  rf.me,
		LastLogIndex: rf.lastLogIndex(),
		LastLogTerm:  rf.lastLogTerm(),
	}
	rf.mu.Unlock()
	for i := range rf.peers {
		if i == rf.me {
			continue
		}
		if rf.role == Candidate {
			go rf.sendRequestVote(i, &reqArg, &RequestVoteReply{})
		}
	}
}

func (rf *Raft) leaderAppendEntries() {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	for i := range rf.peers {
		if i == rf.me {
			continue
		}
		if rf.role == Leader {
			appReply := AppendEntriesReply{}
			appArg := AppendEntriesArgs{
				Term:         rf.currentTerm,
				LeaderId:     rf.me,
				LeaderCommit: rf.commitIndex,
				PrevLogIndex: rf.nextIndex[i] - 1,
			}
			// if server has logs
			if appArg.PrevLogIndex >= 0 {
				appArg.PrevLogTerm = rf.log[appArg.PrevLogIndex].Term
			} else {
				appArg.PrevLogTerm = 0
			}
			// if leader has new logs, especially kick out empty leader
			// lastLogIndex is -1 at the beginning
			if rf.nextIndex[i] <= rf.lastLogIndex() && rf.nextIndex[i] >= 0 {
				// add all logs in and after nextIndex
				rf.debugLogger2B(fmt.Sprintf("server: %v, nextIndex: %v, lastlogIndex: %v", i, rf.nextIndex[i], rf.lastLogIndex()))
				appArg.Entries = rf.log[rf.nextIndex[i]:]
			}
			go rf.sendAppendEntries(i, &appArg, &appReply)
		}
	}
}

func (rf *Raft) applyLogs() {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	for i := rf.lastApplied + 1; i <= rf.commitIndex; i++ {
		rf.ApplyCh <- ApplyMsg{
			Command:      rf.log[i].Command,
			CommandIndex: i,
			CommandValid: true,
		}
	}
	rf.lastApplied = rf.commitIndex
}

func (rf *Raft) lastLogIndex() int {
	return rf.log[len(rf.log)-1].Index
}

func (rf *Raft) lastLogTerm() int {
	return rf.log[len(rf.log)-1].Term
}

func (rf *Raft) debugLogger(s string) {
	var r string
	if rf.role == Leader {
		r = "    Leader"
	} else if rf.role == Follower {
		r = "  Follower"
	} else {
		r = " Candidate"
	}
	DPrintf("Raft ID:%v, role:%v, term:%v, :: %v", rf.me, r, rf.currentTerm, s)
	r = "2"
	trashCan(r)
}

func trashCan(r string) string {
	return r
}

func (rf *Raft) debugLogger2B(s string) {
	var r string
	if rf.role == Leader {
		r = "    Leader"
	} else if rf.role == Follower {
		r = "  Follower"
	} else {
		r = " Candidate"
	}
	DPrintf("Raft ID:%v, role:%v, term:%v, :: %v", rf.me, r, rf.currentTerm, s)
}
